#include "ScoreMatrix.h"
#include "Score.h"
#include "CNF.h"

using namespace std;

#define DUMMY -1
#define ATOM_TYPE_CA 0
#define ATOM_TYPE_CB 1

#define MAX_COL 500
#define _MAX_PAIRWISE 10

class PNN1;

class SEQUENCE
{
public:
  SEQUENCE(int len, PNN1* pModel, string name);
  virtual ~SEQUENCE();

  PNN1* m_pModel;

  static Score Gate(Score sum)
  { return (Score)1.0/(1.0+exp(-(double)sum)); }

  string p_name;
  int length_seq;
  int nDataPoints;

  double Rg;
  int length_contact;

  int* obs_label;
  Score **obs_feature;

  ScoreMatrix* distances;
  
  vector<int> gaps;

  ScoreMatrix* arrVi;

  Array3D<double> *obs_pairwise_features;

public:
  void ComputeVi();
  void CalcPartition();
  virtual Score Predict(){};
  Score MAP();
  Score Obj();

  virtual int makeFeatures(){return 0;};

  virtual void ComputeGates(){};
  virtual Score ComputeGradient(bool bCalculateGate) { return 0;};
  
  virtual void ComputeTestAccuracy(){};

protected:
  int currFeatureIdx;

  int dim_features, window_size, dim_one_pos, num_states, num_gates, 
      dim_pairwise_features, num_labels;
  Score* pFeatures;

  ScoreMatrix *output_center;

  ScoreMatrix* gates;

  int s_start, s_end;

  int *predicted_label;
  Score* Partition;

  double* obj_weights;
};  

class SingletonSequence: public SEQUENCE
{
public:
  SingletonSequence(int len, PNN1* pModel, string name);
  ~SingletonSequence();

  int makeFeatures();

  Score Predict();
  void ComputeGates();
  Score ComputeGradient(bool bCalculateGate);
  
  void ComputeTestAccuracy();

protected:
  ScoreMatrix* _features;
};

class PairwiseSequence: public SEQUENCE
{
public:
  PairwiseSequence(int len, PNN1* pModel, string name);
  ~PairwiseSequence();

  virtual int makeFeatures();
  virtual Score getFeatures(int pos, int i);
  virtual Score Predict();

  virtual void ComputeGates();
  Score ComputeGradient(bool bCalculateGate);
  
  void ComputeTestAccuracy();

protected:
  int *IdxMap, *rIdxMap, *centerMap, *IdxMapCofactor, *IdxMapCenter;
  ScoreMatrix *output_cofactor;
  ScoreMatrix *pairwise_features;
  int contact_DIST;

  int *_features;
};

class DistancePotentialSequence: public PairwiseSequence
{
public:
  DistancePotentialSequence(int len, PNN1* pModel, string name);

  int makeFeatures();
  Score getFeatures(int pos, int i);
  //virtual void ComputeGates();
};

class ContactNumberSequence: public PairwiseSequence
{
public:
  ContactNumberSequence(int len, PNN1* pModel, string name);

  virtual int makeFeatures();
  virtual Score getFeatures(int pos, int i);
  //virtual void ComputeGates();
  virtual Score Predict();

private:
  int* obs_cnum;
};

class _LBFGS: public Minimizer
{
public:
  _LBFGS(PNN1* pModel): Minimizer(false) {m_pModel = pModel;};

  PNN1* m_pModel;

  double ComputeGradient(int iteration, vector<double> &g, 
                         const vector<double> &x, bool bCalculateGate, 
                         int num_consec_small_steps);
  double ComputeFunction(int iteration, const vector<double> &x);
  void Report(const vector<double> &theta, int iteration, 
              double objective, double step_length);
  void Report(const string &s);
};

class PNN1
{
public:
  int num_labels;
  int num_states;
  int num_data;
  int num_tst;
  int num_gates;
  int dim_one_pos;
  int dim_pairwise_features;
  int dim_features;
  int window_size;
  int num_params;
  int totalPos, totalCorrect,totalTriangle,legalTriangle;

  enum MODEL {singleton=0, pairwise=1, DistancePotential=2, contact_number=3} model;
  string model_file;

  //int ct;

  vector<SEQUENCE*> trainData,testData;

  int nnz_param;

  int iGateStart, iFeatureStart;
  double* grad;
  double* weights;
  double* grad_sum;
  double* reg;

  // get the transition prob, num_states( for DUMMY ) + num_states*num_states
  inline Score GetWeightT(int leftState, int currLabel) { 
    currWeightTidx = num_labels * (leftState +1)+ currLabel; 
    return weights[currWeightTidx]; 
  }
  
  inline Score GetWeightTNextLabel()
  { currWeightTidx++; return weights[currWeightTidx];}

  // get the label-based weight, num_states*num_gates
  inline Score* GetWeightLAddr(int currState, int gate)
  { currWeightLidx=iGateStart+ currState*num_gates + gate; return &weights[currWeightLidx];}

  inline Score GetWeightL(int currState, int gate){
    currWeightLidx=iGateStart+ currState*num_gates + gate;
    return weights[currWeightLidx];
  }

  inline Score GetWeightLNextGate()
  { currWeightLidx++; return weights[currWeightLidx]; }

  // get the Gate weight
  inline Score* GetWeightGAddr(int gate, int dim)
  { currWeightGidx=iFeatureStart+gate*dim_features + dim; return &weights[currWeightGidx]; }

  inline Score GetWeightG(int gate, int dim) {
    currWeightGidx=iFeatureStart+gate*dim_features + dim; return weights[currWeightGidx];
  }

  inline Score GetWeightGNextDim()
  { currWeightGidx++; return weights[currWeightGidx]; }

  void SetSeed();
  void SetParameters(int model, int w_size, int n_labels, int n_gates,
                     int n_local, int n_pairwise);
  void Initialize(int model, string model_dir, int w_size, int n_labels,
                  int n_gates, string columns, string pairwise_columns,
                  string input_f, string split_file);
  void LoadData(string input, string split_file);
  void LoadDataSVMFormat(string input);
  void CopyWeightsOut(vector<double>& w);
  void CopyWeightsIn(const vector<double>& w);
  double CopyGradsOut( vector<double>& g, int& nonzeroDirection, 
                       int num_consec_small_steps);
  double RegularizeWeights();
  void RegularizeGrads();

  void convertModel(string oldModel, string newModel);

  int column_flags[MAX_COL];
  int pairwise_column_flags[MAX_COL];
  static int getColumns(string columns, int* column_flags);

private:
  int currWeightGidx, currWeightTidx, currWeightLidx;
};
